﻿#include <windows.h>
#include <stdio.h>
#include <detours.h>
#include <string>

std::wstring GetExecutableDirectory() {
    wchar_t buffer[MAX_PATH] = { 0 };
    DWORD length = GetModuleFileNameW(NULL, buffer, MAX_PATH);

    if (length == 0) {
        return L"";
    }

    // 移除文件名部分（找到最后一个反斜杠）
    for (DWORD i = length; i > 0; --i) {
        if (buffer[i] == L'\\' || buffer[i] == L'/') {
            buffer[i + 1] = L'\0'; // 截断字符串
            break;
        }
    }

    return std::wstring(buffer);
}

int wmain(int argc, wchar_t* argv[]) {
    BOOL result = false;

    if (argc < 2) {
        wprintf(L"Usage: ikget_win <compile_command>\n");
        return 1;
    }

    // 设置输出文件路径（当前目录的 compile_log.txt）
    wchar_t workingDir[MAX_PATH];
    GetCurrentDirectoryW(MAX_PATH, workingDir);
    wprintf(L"[Working Dir]%s\n", workingDir);
    SetEnvironmentVariableW(L"IKGET_WORKING_DIR", workingDir);


    // 构造命令行（仅对含空格的参数加引号）
    wchar_t cmdLine[4096] = { 0 };
    for (int i = 1; i < argc; ++i) {
        if (wcschr(argv[i], L' ')) {
            wcscat_s(cmdLine, L"\"");
            wcscat_s(cmdLine, argv[i]);
            wcscat_s(cmdLine, L"\" ");
        }
        else {
            wcscat_s(cmdLine, argv[i]);
            wcscat_s(cmdLine, L" ");
        }
    }

    // 使用 Detours 启动目标进程并注入 DLL
    STARTUPINFOW si = { sizeof(si) };
    PROCESS_INFORMATION pi;
    DWORD flags = CREATE_SUSPENDED ;
    //DWORD flags = CREATE_NEW_CONSOLE | CREATE_SUSPENDED;

    // 注入 Hook DLL（需绝对路径）
    // 关键调用：创建进程并注入 DLL
    std::wstring dllPath = GetExecutableDirectory();
    SetEnvironmentVariableW(L"IKGET_DLL_DIR", dllPath.c_str());
    //wprintf(L"dll dir: %s\n", dllPath.c_str());
    dllPath += L"\\ikgethook.dll";
    
    char dllPathA[MAX_PATH];
    WideCharToMultiByte(CP_ACP, 0, dllPath.c_str(), -1, dllPathA, MAX_PATH, NULL, NULL);
    //wprintf(L"cmdLine: %s\n", cmdLine);;
    SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
    result = DetourCreateProcessWithDllW(
        NULL,               // 应用程序名（从命令行解析）
        cmdLine,            // 完整命令行
        NULL, NULL, FALSE, flags, NULL, NULL,
        &si, &pi,
        dllPathA,            // 要注入的 DLL
        NULL                // 原 CreateProcess 函数
    );
    if (result) {
        // 恢复进程执行
        ResumeThread(pi.hThread);
        Sleep(100);
        WaitForSingleObject(pi.hProcess, INFINITE);    
    }
    else{
        wprintf(L"call CreateProcess fail.\n");
    }
    
    SetEnvironmentVariableW(L"IKGET_WORKING_DIR", NULL);
    SetEnvironmentVariableW(L"IKGET_DLL_DIR", NULL);
    SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
    // 清理句柄
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
    return 0;
}
